package midend.IR_Optimizer;

import midend.IR_Value.*;
import midend.IR_Instruction.*;
import midend.IR_Type.*;

import java.util.*;

import static midend.IR_Builder.IR_Atendwithprenext;


public class Tail_Recursion_Eli implements IR_Opt{
    /*
    优化示例：
    优化前：
    define i32 @factorial(i32 %n, i32 %acc) {
    entry:
      %cmp = icmp eq i32 %n, 0
      br i1 %cmp, label %exit, label %recursive

    recursive:
      %n_minus_1 = sub i32 %n, 1
      %new_acc = mul i32 %n, %acc
      %result = call i32 @factorial(i32 %n_minus_1, i32 %new_acc)
      ret i32 %result

    exit:
      ret i32 %acc

    优化后：
    define i32 @factorial(i32 %n, i32 %acc) {
    entry:
      br label %phiBlock

    phiBlock:
      %phi_n = phi i32 [%n, %entry], [%n_minus_1, %recursive]
      %phi_acc = phi i32 [%acc, %entry], [%new_acc, %recursive]
      %cmp = icmp eq i32 %phi_n, 0
      br i1 %cmp, label %exit, label %recursive

    recursive:
      %n_minus_1 = sub i32 %phi_n, 1
      %new_acc = mul i32 %phi_n, %phi_acc
      br label %phiBlock

    exit:
      ret i32 %phi_acc
    }
}*/
    @Override
    public void Optimize(IR_FuncBlock function) {
        //判断是否含有尾递归
        if (!Tail_Recursion(function)) {
            return;
        }
        if (function.get_params().size() > 50){
            return;
        }
        //新建一个phi指令块实现对参数的选择
        IR_BaseBlock phi_Block = create_Phi_Block(function);
        Map<IR_Value, Phi_Instruction> phi_Instructions = create_Phi_Instructions(function, phi_Block);
        for (IR_BaseBlock bb : function.get_bb_List()) {
            if (is_Tail_Recursion(function, bb)) {
                replace_Tail_Recursion(function, bb, phi_Block, phi_Instructions);
            }
        }

        //function.updateInstructionwithBlock();
    }

    private boolean Tail_Recursion(IR_FuncBlock function) {
        for (IR_BaseBlock bb : function.get_bb_List()) {
            if (is_Tail_Recursion(function, bb)) {
                return true;
            }
        }
        return false;
    }

    private boolean is_Tail_Recursion(IR_FuncBlock function, IR_BaseBlock bb) {
        List<Instruction> instructions = bb.get_Instruction_List();
        if (instructions.size() <= 1){
            return false;
        }
        Instruction lastInstruction = instructions.get(instructions.size() - 1);
        Instruction lastButOneInstruction = instructions.get(instructions.size() - 2);
        // 递归调用的结果要返回
//        %result = call i32 @factorial(i32 %n_minus_1, i32 %new_acc)
//        ret i32 %result
        if (lastInstruction instanceof Ret_Instruction ret_Inst
                && lastButOneInstruction instanceof Call_Instruction call_Inst) {
            return call_Inst.get_funcBlock().get_FuncName().equals(function.get_FuncName())
                    && ret_Inst.get_op().get(0).get_ValueName().equals(call_Inst.get_op().get(0).get_ValueName());
        }
        return false;
    }

    //在函数入口块后插入一个新的基本块（phiBlock）
    //将原入口块的指令移动到phiBlock
    //在原入口块末尾添加跳转到phiBlock的分支指令
    private IR_BaseBlock create_Phi_Block(IR_FuncBlock function) {
        IR_BaseBlock entry = function.get_bb_List().get(0);
        List<IR_BaseBlock> entry_successors = entry.getNextBlock();
        IR_BaseBlock phi_Block = new IR_BaseBlock("phiBlock");
        function.addbb(phi_Block);
        IR_Atendwithprenext(entry,phi_Block);
        for(IR_BaseBlock successor : entry_successors) {
            IR_Atendwithprenext(phi_Block, successor);
        }
        List<Instruction> entry_Instructions = entry.get_Instruction_List();
        List<Instruction> phi_Instructions = new ArrayList<>(entry_Instructions);
        for (Instruction instruction : phi_Instructions) {
            phi_Block.add_InstructionList(instruction);
        }
        entry_Instructions.clear();
        entry.add_InstructionList(new Branch_Instruction(Collections.singletonList(phi_Block), entry));
        return phi_Block;
    }

    private Map<IR_Value, Phi_Instruction> create_Phi_Instructions(IR_FuncBlock function, IR_BaseBlock phi_Block) {
        Map<IR_Value, Phi_Instruction> phi_Instructions = new HashMap<>();
        List<IR_Value> params = function.get_params();
        for (IR_Value param : params) {
            IR_Type type = param.get_Type();
            IR_Value resRegister = new IR_VirtualReg_Value("phi", type);
            Phi_Instruction phi_Inst = new Phi_Instruction(Collections.singletonList(resRegister), phi_Block);
            phi_Block.add_InstructionList_at_Start(phi_Inst);
            phi_Inst.set_ComingValue(param, function.get_bb_List().get(0));
            phi_Instructions.put(param, phi_Inst);
        }
        return phi_Instructions;
    }

    private void replace_Tail_Recursion(IR_FuncBlock function, IR_BaseBlock bb, IR_BaseBlock phi_Block, Map<IR_Value, Phi_Instruction> phi_Instructions) {
        //替换尾递归块的call和ret为br
        List<Instruction> instructions = bb.get_Instruction_List();
        Call_Instruction call_Inst = (Call_Instruction) instructions.get(instructions.size() - 2);
        Ret_Instruction ret_Inst = (Ret_Instruction) instructions.get(instructions.size() - 1);
        instructions.remove(call_Inst);
        instructions.remove(ret_Inst);
        bb.add_InstructionList(new Branch_Instruction(Collections.singletonList(phi_Block), bb)); // 在bb中增加一条跳转到phi_block的指令
        List<IR_Value> params = function.get_params();
        for (int i = 0; i < params.size(); i++){
            IR_Value param = function.get_param(i);
            IR_Value args = call_Inst.get_params().get(i);
            phi_Instructions.get(param).set_ComingValue(args, bb);

            IR_Value new_Value = phi_Instructions.get(param).get_op().get(0);
            //将形参在函数中的引用点替换为Phi指令的结果
            function.replaceUseInFunc(param, new_Value);
        }
    }
}